from __future__ import absolute_import
import logging
import os
import time
import re

from os import mkdir
from os.path import join, abspath, exists #, dirname
#from subprocess import check_call

from forest.api import Command, get_tool

cpfile = get_tool('system.copy')
mvfile = get_tool('system.move')
mkdir = get_tool('system.mkdir')
chroot = get_tool('system.chroot')
execute_tool = execute = get_tool('system.execute')
working_directory = get_tool('system.cd')

logger = logging.getLogger('forest.commands.BuildLiveImage')

class BuildLiveImage(Command):
    ''' Builds a live image

        Args

            disklabel - string - label stored in .discinfo

            rootlabel - string - Label of the live root filesystem

            isolabel - string - Label of the iso

            isolinuxcfg - string - relative or absolute path of isolinux.cfg to include

            syslinux_source - relative or absolute path to where syslinux binaries can be found (eg isolinux.bin, vesamenu)
    '''

    __register__ = 'commands.live.BuildLiveImage'

    def run(self, config):
        self._config = config
        # create a directory for the iso
        Command.SIGNAL_COMMAND_PROGRESS(self, level=logging.INFO, message='Building live image...')

        #install_root = config.get_install_root()

        self.buildisolinux(config)
        self.copysyslinux()
        self.creatediscinfo()
        self.buildliveimage(config)
        self.buildiso()

    def get_iso_root(self):
        return os.path.join(self._config.get_build_root(), 'isoroot')
        
    def join_iso_root(self, *args):
        return os.path.join(self._config.get_build_root(), 'isoroot', *args)
        
    @staticmethod
    def get_size(start_path = '.'):
        # TODO: ignore proc and sys and whatnot
        total_size = 0
        for dirpath, dirnames, filenames in os.walk(start_path):
            for f in filenames:
                fp = os.path.join(dirpath, f)
                if os.path.exists(fp):
                    total_size += os.path.getsize(fp)
        return total_size


    def build_fs_image(self, root, dest, label):
        # package installtree into a ext4fs into a squash

        installtree = root

        # 2GB sparse file compresses down to nothin'
        fssize = 2 * (1024*1024*1024)

        Command.SIGNAL_COMMAND_PROGRESS(self, level=logging.INFO, message='Calculating install tree size...')
        installtreesize = self.get_size(root)

        Command.SIGNAL_COMMAND_PROGRESS(self, level=logging.INFO, message='%s Megabytes' % (installtreesize))
        fssize = 10 * (1024*1024) + installtreesize

        get_tool('image.mkext4image')(installtree, dest, label=label, size=fssize)
        # squash the live rootfs and clean up workdir
        #imgutils.mksquashfs(squashfsdir, squashfsimage, "xz", [])

    def buildisolinux(self, config):
        #install_root = config.get_install_root()

        #configfilepath = dirname(config.get_filename())

        isolinuxcfgfile = config.get_defined_file(self.args.get('isolinuxcfg'))

        Command.SIGNAL_COMMAND_PROGRESS(self, level=logging.INFO, message='Copying isolinux.cfg %s...'% isolinuxcfgfile)

        #TODO - also build a dracut image
        mkdir(self.get_iso_root())
        mkdir(self.join_iso_root('isolinux'))
        
        isolinuxfile = open(isolinuxcfgfile).read().replace('$architecture',config.get('yum.architecture'))
        get_tool('files.write')(self.join_iso_root('isolinux', 'isolinux.cfg'), isolinuxfile)

    def copysyslinux(self):
        Command.SIGNAL_COMMAND_PROGRESS(self, level=logging.INFO, message='Installing isolinux...')

        # create a isolinux directory and structures
        isolinuxdest = self.join_iso_root('isolinux')

        if not exists(isolinuxdest):
            mkdir(isolinuxdest)

        # isolinux/isolinux.bin and vesamenu.c32 from isolinux_source
        syslinux_source = self.args.get('syslinux_source')
        for fn in ['isolinux.bin', 'vesamenu.c32']:
            fnpath = join(syslinux_source, fn)
            cpfile(fnpath, self.join_iso_root('isolinux', fn))

    def creatediscinfo(self):
        discinfo_lst = ['%s\n' % time.time(),
            '%s\n' % self.args.get('disklabel'),
            '%s' % self.args.get('arch')]

        get_tool('files.write')(self.join_iso_root('.discinfo'), ''.join(discinfo_lst), mode='w+')


    def buildliveimage(self, config, destination=None):
        install_root = config.get_install_root()
        #tempdir = config.get('tempdir')
        ext3fsimage = join('build', 'ext3fs.img')

        Command.SIGNAL_COMMAND_PROGRESS(self, None, logging.INFO, 'Building filesystem image...')
        # build the fs image
        self.build_fs_image(install_root, ext3fsimage, self.args.get('rootlabel','root'))

        # copy it to the live
        mkdir(self.join_iso_root('Live'))

        if destination is None:
            destination = self.join_iso_root('Live')
            
        mvfile(ext3fsimage, destination)

        return ext3fsimage

    def buildiso(self):
        Command.SIGNAL_COMMAND_PROGRESS(self, level=logging.INFO, message='Building iso...')
        dest = abspath(self.args.get('destination'))
        with working_directory(self.get_iso_root()):
            get_tool('image.mkisofs')(dest=dest, bootloader="isolinux/isolinux.bin", label=self.args.get('isolabel'), bootcat="isolinux/boot.cat")

    def findkernels(self, root="/", kdir="boot"):
        # To find possible flavors, awk '/BuildKernel/ { print $4 }' kernel.spec
        flavors = ('debug', 'PAE', 'PAEdebug', 'smp', 'xen')
        kre = re.compile(r"vmlinuz-(?P<version>.+?\.(?P<arch>[a-z0-9_]+)"
                         r"(\.(?P<flavor>{0}))?)$".format("|".join(flavors)))
        kernels = []
        
        if not os.path.exists(join(root, kdir)):
            return kernels
        
        for f in os.listdir(join(root, kdir)):
            match = kre.match(f)
            if match:
                kernel = {'path':join(root, kdir, f)}
                kernel.update(match.groupdict()) # sets version, arch, flavor
                kernels.append(kernel)
                Command.SIGNAL_COMMAND_PROGRESS(self, level=logging.INFO, message='Found kernel %s' % kernel)


        # look for associated initrd/initramfs
        for kernel in kernels:
            # NOTE: if both exist, the last one found will win
            for imgname in ("initrd", "initramfs"):
                i = kernel.get('path').replace("vmlinuz", imgname, 1) + ".img"
                if os.path.exists(join(root, i)):
                    kernel['initrd_path'] = i

        return kernels

    def build_initrd(self, root, path):
        dracut = ["/sbin/dracut",  "--nomdadmconf", "--nolvmconf",'--force']
        kernels = self.findkernels(root=root)

        Command.SIGNAL_COMMAND_PROGRESS(self, level=logging.INFO, message="Found kernels: %s" % kernels)

        kernel = kernels[0]

        Command.SIGNAL_COMMAND_PROGRESS(self, level=logging.INFO, message="Rebuilding initrd to %s. Using kernel %s" % (path, kernel))
        #if backup:
        #    initrd = join(root, path)
        #    os.rename(initrd, initrd + backup)
        with chroot(root):
            args = dracut + [path, kernel.get('version')]
            for line in get_tool('system.poll')(args):
                print line

        #os.unlink(join(root,"/proc/modules"))
        return kernel

    def generate_dracut_initrd(self, config, path):
        ''' Generates a dracut initrd at path '''
        return self.build_initrd(config.get_install_root(), path)

#import re
#import os
#from subprocess import check_call
#from os.path import join
#
#def rebuild_initrds(self, root, add_args=[], backup=""):
#    '''Rebuild all the initrds in the tree. If backup is specified, each
#    initrd will be renamed with backup as a suffix before rebuilding.
#    If backup is empty, the existing initrd files will be overwritten.'''
#
#    dracut = ["/sbin/dracut", "--noprefix", "--nomdadmconf", "--nolvmconf"] + add_args
#
#    if not backup:
#        dracut.append("--force")
#
#    # Hush some dracut warnings. TODO: bind-mount proc in place?
#    open(join(root,"/proc/modules"),"w")
#
#    kernels = findkernels(root=root)
#
#    for kernel in self.kernels:
#        logger.info("rebuilding %s", kernel.get('initrd_path'))
#        if backup:
#            initrd = join(root, kernel.get('initrd_path'))
#            os.rename(initrd, initrd + backup)
#        check_call(["chroot", root] + dracut + [kernel.initrd.path, kernel.version])
#    os.unlink(join(root,"/proc/modules"))
